# 实现基本运算：
import torch
# 搭建网络结构：
import torch.nn as nn
# 实现前馈运算：
import torch.nn.functional as f

# file:CNN.py
class ConvNet(nn.Module):
    # """编写一个卷积神经网络类"""
	def __init__(self):
		""" 初始化网络,将网络需要的模块拼凑出来。 """
		super(ConvNet, self).__init__()
		# 卷积层:
		self.conv1 = nn.Conv2d(1, 6, 5, padding=2)
		self.conv2 = nn.Conv2d(6, 16, 5, padding=2)
		# 最大池化处理:
		self.pooling = nn.MaxPool2d(2, 2)
		# 全连接层：
		self.fc1 = nn.Linear(16*7*7, 512)
		self.fc2 = nn.Linear(512, 10)

	def forward(self, x):
		"""前馈函数"""
		x = f.relu(self.conv1(x)) # = [b, 6, 28, 28]
		x = self.pooling(x)       # = [b, 6, 14, 14]
		x = f.relu(self.conv2(x)) # = [b, 16, 14, 14]
		x = self.pooling(x)		  # = [b, 16, 7, 7]
		x = x.view(x.shape[0], -1)# = [b, 16 * 7 * 7]
		x = f.relu(self.fc1(x))
		x = self.fc2(x)
		output = f.log_softmax(x, dim=1)
		return output

def cal_correction(output, target):
    _,predicted = torch.max(output, 1)  # Get the index of the max log-probability
    correct = predicted.eq(target.data.view(-1))  # Flatten target if needed
    accuracy = (correct.sum().item() / len(correct))*100
    return accuracy